from copy import deepcopy
import numpy as np
from typing import Union
from . import pregel as core
from . import ADMM_update
from .ADMM_update import terminal_args

"""
已确定graphscope的pregel调不了求解器，这部分就转向手搓一个pregel

pregel的python手搓版：https://github.com/mnielsen/Pregel

pregel的C++手搓版：https://github.com/yaobaiwei/PregelPlus
https://www.cse.cuhk.edu.hk/pregelplus/index.html

提供了pregel的代数编程框架：https://github.com/Algebraic-Programming/ALP
"""


# 这个可能没用了
class terminal_message_manager:
    messages: dict[str, Union[list, int, float]]

    def register_message_aggregator(self, key: str):
        val = self.messages.get(key, default=None)
        if val is None or len(val) == 0:
            self.messages.update({key, []})
        else:
            KeyError(f'key "{key}" has already existed')

    def overwrite_message(self, key: str, value):
        self.messages[key] = value

    def get_message(self, key: str):
        return self.messages.get(key, default=None)


def context(**kwargs) -> dict:
    c = {
        "acline": None,
        "storage": None,
        "upgrid": None,
        "acload": None,
        "rou": 1e-2,
        "T": 24,
        "logger": None,
        "debug": False,
        "interval_seconds": 900,
    }
    c.update(kwargs)
    return c


def is_update_converge(
    z_adj_terminals: terminal_args, x_adj_terminals: terminal_args, esp
) -> list[bool]:
    converged = [True for k in range(0, len(x_adj_terminals))]
    for k in range(0, len(x_adj_terminals)):
        if x_adj_terminals[k].var_status("P") == x_adj_terminals[k].args_status.INUSE:
            diff = np.abs(z_adj_terminals[k].P - x_adj_terminals[k].P)
            if not np.max(diff) <= esp:
                converged[k] = False
                continue

        if x_adj_terminals[k].var_status("Q") == x_adj_terminals[k].args_status.INUSE:
            diff = np.abs(z_adj_terminals[k].Q - x_adj_terminals[k].Q)
            if not np.max(diff) <= esp:
                converged[k] = False
                continue

        if x_adj_terminals[k].var_status("V") == x_adj_terminals[k].args_status.INUSE:
            diff = np.abs(z_adj_terminals[k].V2 - x_adj_terminals[k].V2)
            if not np.max(diff) <= esp:
                converged[k] = False
                continue
    return converged


class Vertex_bus(core.Vertex):
    class value_type:
        name: str
        history_terminals: dict

        def __init__(self, name):
            self.name = name
            self.history_terminals = {}
            self.terminal_list = []

    def init(self, context: dict):
        # 这个问题是双向图，出边也是入边
        v = self.getValue()
        for neighbor in self.get_out_vertices():
            # h for handle, m for message
            v.history_terminals[neighbor.getid()] = {"h": neighbor, "m": None}
            v.terminal_list.append(neighbor.getid())
        self.setValue(v)
        self.vote_to_halt()

    # 所有信息交互应该改成用handle，使用vid实现
    def compute(self, messages, context: dict):
        if messages is None or len(messages) == 0:
            raise NotImplementedError("没有消息不能唤醒or最后忘了halt了")

        # merge
        terminal_messages = dict([(src.getid(), msg) for src, msg in messages])
        adj_terminals = []
        val = self.getValue()
        Nt = len(val.terminal_list)
        for k in range(0, Nt):
            if terminal_messages.__contains__(val.terminal_list[k]):
                adj_terminals.append(terminal_messages[val.terminal_list[k]])
            else:
                adj_terminals.append(val.history_terminals[val.terminal_list[k]]["m"])

        # update
        up_adj_terminals = ADMM_update.node_update(
            adj_terminals, context["rou"], context["T"]
        )

        # 一致性收敛判断
        consensus_list = is_update_converge(
            up_adj_terminals,
            adj_terminals,
            context["tolerance"] * np.sqrt(context["T"] * Nt),
        )

        # 步进收敛判断
        if self.superstep > 1:
            step_list = is_update_converge(
                up_adj_terminals,
                [val.history_terminals[val.terminal_list[k]]["m"] for k in range(0, Nt)],
                context["tolerance"] * np.sqrt(context["T"] * Nt),
            )
        else:
            step_list = [False for k in range(0, Nt)]

        # store and send
        for k in range(0, Nt):
            if consensus_list[k] is False or step_list[k] is False:
                val.history_terminals[val.terminal_list[k]]["m"] = up_adj_terminals[k]
                self.send(
                    val.history_terminals[val.terminal_list[k]]["h"], up_adj_terminals[k]
                )
        self.setValue(val)
        # logging
        if (context["logger"] is not None) and (context["debug"] is True):
            self.logging(context["logger"])
        self.vote_to_halt()

    def logging(self, logger):
        value = self.value
        logger.debug(f"Device acbus {value.name} is active at superstep {self.superstep}:")


class bus_builder:
    value: Vertex_bus.value_type
    vid: int
    out_vertices_id: Union[list, tuple]

    def __init__(self, vid, name: str, out_vertices_id: Union[list, tuple]):
        self.vid = vid
        self.value = Vertex_bus.value_type(name)
        self.out_vertices_id = out_vertices_id

    def get_out_vertices_id(self) -> Union[list, tuple]:
        return self.out_vertices_id

    def build(self) -> Vertex_bus:
        return Vertex_bus(self.vid, self.value)


class Vertex_acline(core.Vertex):
    class value_type:
        name: str
        history_terminal1: terminal_args
        history_terminal2: terminal_args

        def __init__(self, name):
            self.name = name


    def logging(self, logger):
        value = self.value
        logger.debug(f"Device acline {value.name} at superstep {self.superstep}:")

    def init(self, context: dict):
        # 拉格朗日乘子设为0，求完给bus直接发消息
        v= self.getValue()
        property_m = context["acline"].query(v.name)
        outv = self.get_out_vertices()
        v.bus1 = outv[0]
        v.bus2 = outv[1]
        self.setValue(v)

        terminal1 = property_m["init_terminal1"]
        self.send(v.bus1, terminal1)

        terminal2 = property_m["init_terminal2"]
        self.send(v.bus2, terminal2)
        self.vote_to_halt()

    def compute(self, messages: list, context: dict):
        # 有新的用新的，没新的用旧的
        skip1 = True
        skip2 = True
        val=self.getValue()
        terminal_messages = dict([(src.getid(),msg) for src, msg in messages])
        if terminal_messages.__contains__(val.bus1.getid()):
            terminal1 = terminal_messages[val.bus1.getid()]
            skip1 = False
        else:
            terminal1 = val.history_terminal1

        if terminal_messages.__contains__(val.bus2.getid()):
            terminal2 = terminal_messages[val.bus2.getid()]
            skip2 = False
        else:
            terminal2 = val.history_terminal2
        if skip1 is True and skip2 is True:
            # 没收到消息
            return
        property_m = context["acline"].query(val.name)
        opt_sol = ADMM_update.acline_update(
            property_m["R"],
            property_m["X"],
            property_m["Vmax_i"],
            property_m["Vmin_j"],
            property_m["Imax"],
            terminal1,
            terminal2,
            context["rou"],
            context["T"],
        )

        # store and send
        terminal1.P = opt_sol["Pi"]
        terminal1.Q = opt_sol["Qi"]
        terminal1.V2 = opt_sol["V2i"]
        terminal1.LP = opt_sol["LPi"]
        terminal1.LQ = opt_sol["LQi"]
        terminal1.LV2 = opt_sol["LV2i"]
        val.history_terminal1 = terminal1

        terminal2.P = opt_sol["Pj"]
        terminal2.Q = opt_sol["Qj"]
        terminal2.V2 = opt_sol["V2j"]
        terminal2.LP = opt_sol["LPj"]
        terminal2.LQ = opt_sol["LQj"]
        terminal2.LV2 = opt_sol["LV2j"]
        val.history_terminal2 = terminal2

        self.send(val.bus1, terminal1)
        self.send(val.bus2, terminal2)

        if (context["logger"] is not None) and (context["debug"] is True):
            self.logging(context["logger"])
        self.vote_to_halt()


class acline_builder:
    value: Vertex_acline.value_type
    vid = None
    out_vertices_id: Union[list, tuple]

    def __init__(self, vid, name: str, out_vertices_id: Union[list, tuple]):
        self.vid = vid
        self.value = Vertex_acline.value_type(name)
        self.out_vertices_id = out_vertices_id

    def get_out_vertices_id(self) -> Union[list, tuple]:
        return self.out_vertices_id

    def build(self) -> Vertex_acline:
        return Vertex_acline(self.vid, self.value)


class Vertex_storage(core.Vertex):
    class value_type:
        name: str
        Pch: np.ndarray
        Pdis: np.ndarray
        Pi: np.ndarray
        state_ch: np.ndarray
        Ee: np.ndarray

        def __init__(self, name):
            self.name = name

    def init(self, context: dict):
        # 根据初始化参数给出当前策略
        property_map = context["storage"].query(self.value.name)
        terminal1 = property_map["init_terminal"]
        self.send(self.get_out_vertices()[0], terminal1)
        self.vote_to_halt()

    def compute(self, messages: list, context: dict):
        # 只会收到一条消息
        if len(messages) > 1:
            raise RuntimeError("[storage]receive unexpected number of messages")
        elif len(messages) == 0:
            self.vote_to_halt()
        # update
        property_map = context["storage"].query(self.value.name)
        terminal1=messages[0][1]
        y = ADMM_update.storage_update(
            PEssMax=property_map["Pmax"],
            eff_ch=property_map["eff_ch"],
            eff_dis=property_map["eff_dis"],
            capacity=property_map["capacity"],
            max_changing_N=property_map["max changing"],
            cost_operating=property_map["operating cost"],
            depreciation_per_change=property_map["changing cost"],
            interval_seconds=context["interval_seconds"],
            terminal=terminal1,
            rou=context["rou"],
            T=context["T"],
        )
        terminal1.P = y["Pi"]
        terminal1.Q = y["Qi"]
        terminal1.LP = y["LPi"]
        terminal1.LQ = y["LQi"]

        # send message
        self.send(messages[0][0], terminal1)

        # logging
        if (context["logger"] is not None) and (context["debug"] is True):
            self.logging(context["logger"])

        self.vote_to_halt()

    def logging(self, logger):
        info = f"Device storage {self.value.name} at superstep {self.superstep}:"
        """
        info += f"Pi:{self.value.Pi}"
        info += f"Ee:{self.value.Ee}"
        info += f"switch state:{self.value.state_ch}"
        """
        logger.debug(info)


class storage_builder:
    value: Vertex_storage.value_type
    vid = None
    out_vertices_id: Union[list, tuple]

    def __init__(self, vid, name: str, out_vertices_id: Union[list, tuple]):
        self.vid = vid
        self.value = Vertex_storage.value_type(name)
        self.out_vertices_id = out_vertices_id

    def get_out_vertices_id(self) -> Union[list, tuple]:
        return self.out_vertices_id

    def build(self) -> Vertex_storage:
        return Vertex_storage(self.vid, self.value)


class Vertex_acsource(core.Vertex):
    class value_type:
        name: str
        Psub: np.ndarray
        Qsub: np.ndarray
        Cupper: np.ndarray

        def __init__(self, name):
            self.name = name
            self.Psub = None
            self.Qsub = None

    def logging(self, logger):
        info = f"Device upper grid {self.value.name} at superstep {self.superstep}:"
        '''
        info += f"Psub:{self.value.Psub}"
        info += f"Cupper:{self.value.Cupper}"
        '''
        logger.debug(info)

    def init(self, context: dict):
        property_map = context["upgrid"].query(self.value.name)
        terminal1 = property_map["init_terminal"]
        self.send(self.get_out_vertices()[0], terminal1)
        self.vote_to_halt

    def compute(self, messages, context: dict):
        # 只会收到一条消息
        if len(messages) > 1:
            raise RuntimeError("[upgrid]receive unexpected number of messages")
        elif len(messages) == 0:
            self.vote_to_halt()
            return 
        # update
        property_map = context["upgrid"].query(self.value.name)
        terminal1= messages[0][1]
        sol = ADMM_update.upper_grid_update(
            property_map["Ct"],
            property_map["Pmax"],
            property_map["Pmin"],
            terminal1,
            context["rou"],
            context["T"],
        )

        # store and send
        terminal1.P = sol["Pi"]
        terminal1.Q = sol["Qi"]
        terminal1.LP = sol["LPi"]
        terminal1.LQ = sol["LQi"]
        terminal1.V2 = np.ones((context["T"], 1), np.float64) * (
            property_map["kV"] ** 2
        )
        self.send(messages[0][0], terminal1)

        # logging
        if (context["logger"] is not None) and (context["debug"] is True):
            self.logging(context["logger"])

        self.vote_to_halt()


class acsource_builder:
    value: Vertex_acsource.value_type
    vid = None
    out_vertices_id: Union[list, tuple]

    def __init__(self, vid, name: str, out_vertices_id: Union[list, tuple]):
        self.vid = vid
        self.value = Vertex_acsource.value_type(name)
        self.out_vertices_id = out_vertices_id

    def get_out_vertices_id(self) -> Union[list, tuple]:
        return self.out_vertices_id

    def build(self) -> Vertex_acsource:
        return Vertex_acsource(self.vid, self.value)


class Vertex_load(core.Vertex):
    class value_type:
        name: str
        Pi: np.ndarray
        Qi: np.ndarray

        def __init__(self, name):
            self.name = name
            self.Pi = None
            self.Qi = None

    def init(self,context: dict):
        property_m = context["acload"].query(self.value.name)
        init_terminal = ADMM_update.load_update(
            property_m["Pt"],
            property_m["Qt"],
            terminal_args(context["T"], ("P", "Q")),
            context["rou"],
            context["T"],
        )
        self.send(self.get_out_vertices()[0], init_terminal)
        self.vote_to_halt()

    def compute(self, messages, context: dict):
        # 只会收到一条消息
        if len(messages) > 1:
            raise RuntimeError("[Load]receive unexpected number of messages")
        elif len(messages)==0:
            self.vote_to_halt()
            return
        # update
        property_map = context["acload"].query(self.value.name)
        update_res = ADMM_update.load_update(
            property_map["Pt"],
            property_map["Qt"],
            messages[0][1],
            context["rou"],
            context["T"],
        )
        # store and send
        self.send(self.get_out_vertices()[0], update_res)

        self.vote_to_halt()


class load_builder:
    value: Vertex_load.value_type
    vid = None
    out_vertices_id: Union[list, tuple]

    def __init__(self, vid, name: str, out_vertices_id: Union[list, tuple]):
        self.vid = vid
        self.value = Vertex_load.value_type(name)
        self.out_vertices_id = out_vertices_id

    def get_out_vertices_id(self) -> Union[list, tuple]:
        return self.out_vertices_id

    def build(self) -> Vertex_load:
        return Vertex_load(self.vid, self.value)
